package tool

import (
	"colly/pkg/hook"
	"context"
	"fmt"
	"github.com/google/uuid"
	"github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"
	"io"
	"os"
	"time"
)

type ElapsedTimeHook struct {
	StartTime time.Time
	Context   context.Context
}

// NewElapsedTimeHook creates a new ElapsedTimeHook.
func NewElapsedTimeHook(ctx context.Context) *ElapsedTimeHook {
	return &ElapsedTimeHook{
		StartTime: time.Now(),
		Context:   ctx,
	}
}

// Levels defines on which log levels this hook would trigger.
func (hook *ElapsedTimeHook) Levels() []logrus.Level {
	return logrus.AllLevels
}

// Fire is called when a log entry is fired.
func (hook *ElapsedTimeHook) Fire(entry *logrus.Entry) error {
	elapsed := time.Since(hook.StartTime)
	entry.Data["elapsed_time"] = elapsed.Seconds()

	// Add trace_id from context if available
	if traceID, ok := hook.Context.Value("trace_id").(string); ok {
		entry.Data["trace_id"] = traceID
	}

	return nil
}

func getLogFileName() string {
	// 获取当前日期
	currentTime := time.Now()
	year, month, day := currentTime.Date()

	// 格式化为 "YYYY-MM-DD" 的字符串
	return fmt.Sprintf("./logs/app-%d-%02d-%02d.log", year, month, day)
}

type CustomFormatter struct {
	logrus.JSONFormatter
}

// Format 格式化日志条目
func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	entry.Message = fmt.Sprintf("%s", entry.Message)
	return f.JSONFormatter.Format(entry)
}

func InitLogger() *logrus.Logger {

	traceID := uuid.New()
	fileLogger := &lumberjack.Logger{
		Filename:   getLogFileName(),
		MaxSize:    50,    // Max size in MB
		MaxBackups: 1,     // Max number of old log files to keep
		MaxAge:     7,     // Max age in days to keep old log files
		Compress:   false, // Compress old log files
	}
	log := logrus.New()
	ctx := context.WithValue(context.Background(), "trace_id", traceID)
	log.SetOutput(io.MultiWriter(os.Stdout, fileLogger))
	//log.SetFormatter(&logrus.JSONFormatter{
	//	DisableTimestamp: false,
	//	TimestampFormat:  "2006-01-02 15:04:05",
	//	FieldMap:         nil,
	//	CallerPrettyfier: nil,
	//})

	log.SetFormatter(&CustomFormatter{
		JSONFormatter: logrus.JSONFormatter{
			DisableTimestamp: false,
			TimestampFormat:  "2006-01-02 15:04:05",
			FieldMap:         nil,
			CallerPrettyfier: nil,
		},
	})

	// Add the custom hook with the context
	log.AddHook(&hook.ContextHook{Context: ctx})
	return log
}
